﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using TomanuExtensions;
using RaytracerLib.Layers.OverlayCorrectors;
using System.Threading.Tasks;
using System.Drawing;
using System.Diagnostics;
using TomanuExtensions.Utils;
using System.Numerics;

namespace RaytracerLib.MathLib
{
    public static class FFT
    {
        #region FFT

        public static Complex[] fft(Complex[] a_ar)
        {
            if (Bits.IsBase2(a_ar.Length))
                return fft_radix2(a_ar, a_ar.Length, 0, 1);
            else
                return fft_slow(a_ar);
        }

        public static Complex[] fft(double[] a_ar)
        {
            if (Bits.IsBase2(a_ar.Length))
                return fft_radix2(a_ar, a_ar.Length, 0, 1);
            else
                return fft_slow(a_ar);
        }

        private static Complex[] fft_radix2(Complex[] a_ar, int a_length, int a_from, int a_step)
        {
            if (a_length == 1)
                return new Complex[] { a_ar[a_from] };

            int len2 = a_length / 2;

            Complex[] odds = fft_radix2(a_ar, len2, a_from, a_step * 2);
            Complex[] evens = fft_radix2(a_ar, len2, a_from + a_step, a_step * 2);

            Complex[] result = new Complex[a_length];

            Complex step = Complex.FromPolarCoordinates(1, -2 * MathExtensions.PI / a_length);
            Complex coeff = Complex.One;

            for (int k = 0; k < result.Length / 2; k++)
            {
                Complex even = evens[k];

                Complex temp = coeff * odds[k];
                result[k] = even + temp;
                result[k + len2] = even - temp;
                coeff *= step;
            }

            return result;
        }

        private static Complex[] fft_radix2(double[] a_ar, int a_length, int a_from, int a_step)
        {
            if (a_length == 1)
                return new Complex[] { new Complex(a_ar[a_from], 0) };

            int len2 = a_length / 2;

            Complex[] odds = fft_radix2(a_ar, len2, a_from, a_step * 2);
            Complex[] evens = fft_radix2(a_ar, len2, a_from + a_step, a_step * 2);

            Complex[] result = new Complex[a_length];

            Complex step = Complex.FromPolarCoordinates(1, -2 * MathExtensions.PI / a_length);
            Complex coeff = Complex.One;

            for (int k = 0; k < result.Length/2; k++)
            {
                Complex even = evens[k];

                Complex temp = coeff * odds[k];
                result[k] = even + temp;
                result[k + len2] = even - temp;
                coeff *= step;
            }

            return result;
        }

        private static Complex[] fft_slow(Complex[] a_ar, Complex[,] a_coeffs = null)
        {
            if (a_coeffs == null)
            {
                Complex[] result = new Complex[a_ar.Length];

                for (int k = 0; k < result.Length; k++)
                {
                    Complex step = Complex.FromPolarCoordinates(1, -2 * MathExtensions.PI * k / a_ar.Length);
                    Complex coeff = Complex.One;

                    for (int n = 0; n < a_ar.Length; n++)
                    {
                        result[k] += a_ar[n] * coeff;
                        coeff *= step;
                    }
                }

                return result;
            }
            else
            {
                Complex[] result = new Complex[a_ar.GetLength(0)];

                for (int k = 0; k < result.Length; k++)
                {
                    for (int n = 0; n < a_ar.Length; n++)
                    {
                        result[k] += a_ar[n] * a_coeffs[k, n];
                    }
                }

                return result;
            }
        }

        private static Complex[] fft_slow(double[] a_ar, Complex[,] a_coeffs = null)
        {
            if (a_coeffs == null)
            {
                Complex[] result = new Complex[a_ar.Length];

                for (int k = 0; k < result.Length; k++)
                {
                    Complex step = Complex.FromPolarCoordinates(1, -2 * MathExtensions.PI * k / a_ar.Length);
                    Complex coeff = Complex.One;

                    for (int n = 0; n < a_ar.Length; n++)
                    {
                        result[k] += a_ar[n] * coeff;
                        coeff *= step;
                    }
                }

                return result;
            }
            else
            {
                Complex[] result = new Complex[a_ar.Length];

                for (int k = 0; k < result.Length; k++)
                {
                    for (int n = 0; n < a_ar.Length; n++)
                    {
                        result[k] += a_ar[n] * a_coeffs[k, n];
                    }
                }

                return result;
            }
        }

        #endregion

        #region FFT2

        public static Complex[,] fft2(double[,] a_in)
        {
            int w = a_in.GetLength(0);
            int h = a_in.GetLength(1);

            Complex[,] temp = new Complex[w, a_in.GetLength(1)];
            Complex[,] coeffs = null;

            if (Bits.IsBase2(w))
                PrecalculateCoeffs(a_in.GetLength(0), w);

            Parallel.For(0, a_in.GetLength(1), (y) =>
            {
                double[] row = new double[w];

                for (int x = 0; x < w; x++)
                    row[x] = a_in[x, y];

                Complex[] frow;
                if (Bits.IsBase2(row.Length))
                    frow = fft_radix2(row, row.Length, 0, 1);
                else
                    frow = fft_slow(row, coeffs);

                for (int x = 0; x < w; x++)
                    temp[x, y] = frow[x];

            });

            if (Bits.IsBase2(h))
                coeffs = PrecalculateCoeffs(a_in.GetLength(1), h);

            Complex[,] result2 = new Complex[w, h];

            Parallel.For(0, w, (x) =>
            {
                Complex[] col = new Complex[h];

                for (int y = 0; y < h; y++)
                    col[y] = temp[x, y];

                Complex[] fcol;
                if (Bits.IsBase2(col.Length))
                    fcol = fft_radix2(col, col.Length, 0, 1);
                else
                    fcol = fft_slow(col, coeffs);

                for (int y = 0; y < h; y++)
                    result2[x, y] = fcol[y];
            });

            return result2;
        }

        public static Complex[,] fft2(Complex[,] a_in)
        {
            int w = a_in.GetLength(0);
            int h = a_in.GetLength(1);

            Complex[,] temp = new Complex[w, a_in.GetLength(1)];
            Complex[,] coeffs = null;

            if (Bits.IsBase2(w))
                PrecalculateCoeffs(a_in.GetLength(0), w);

            Parallel.For(0, a_in.GetLength(1), (y) =>
            {
                Complex[] row = new Complex[w];

                for (int x = 0; x < w; x++)
                    row[x] = a_in[x, y];

                Complex[] frow;
                if (Bits.IsBase2(row.Length))
                    frow = fft_radix2(row, row.Length, 0, 1);
                else
                    frow = fft_slow(row, coeffs);

                for (int x = 0; x < w; x++)
                    temp[x, y] = frow[x];

            });

            if (Bits.IsBase2(h))
                coeffs = PrecalculateCoeffs(a_in.GetLength(1), h);

            Complex[,] result2 = new Complex[w, h];

            Parallel.For(0, w, (x) =>
            {
                Complex[] col = new Complex[h];

                for (int y = 0; y < h; y++)
                    col[y] = temp[x, y];

                Complex[] fcol;
                if (Bits.IsBase2(col.Length))
                    fcol = fft_radix2(col, col.Length, 0, 1);
                else
                    fcol = fft_slow(col, coeffs);

                for (int y = 0; y < h; y++)
                    result2[x, y] = fcol[y];
            });

            return result2;
        }

        private static Complex[,] PrecalculateCoeffs(int a_in_length, int a_out_length)
        {
            Complex[,] result = new Complex[a_out_length, a_in_length];

            for (int k = 0; k < a_out_length; k++)
            {
                Complex step = Complex.FromPolarCoordinates(1, -2 * MathExtensions.PI * k / a_in_length);
                Complex coeff = Complex.One;

                for (int n = 0; n < a_in_length; n++)
                {
                    result[k, n] = coeff;
                    coeff *= step;
                }
            }

            return result;
        }

        #endregion

        #region FFT SHIFT

        public static T[] fftshift<T>(T[] a_ar)
        {
            int delta = (a_ar.Length / 2.0).Floor();
            return generic_shift(a_ar, delta);
        }

        public static T[] ifftshift<T>(T[] a_ar)
        {
            int delta = (a_ar.Length / 2.0).Ceiling();
            return generic_shift(a_ar, delta);
        }

        private static T[] generic_shift<T>(T[] a_ar, int delta)
        {
            T[] result = new T[a_ar.Length];

            for (int i = 0; i < a_ar.Length; i++)
            {
                int ci = OverlapOverlayCorrector.Correct_XY(i + delta, a_ar.Length);
                result[ci] = a_ar[i];
            }

            return result;
        }

        public static T[,] fftshift<T>(T[,] a_ar)
        {
            int deltax = (a_ar.GetLength(0) / 2.0).Floor();
            int deltay = (a_ar.GetLength(1) / 2.0).Floor();
            return generic_shift(a_ar, deltax, deltay);
        }

        public static T[,] ifftshift<T>(T[,] a_ar)
        {
            int deltax = (a_ar.GetLength(0) / 2.0).Ceiling();
            int deltay = (a_ar.GetLength(1) / 2.0).Ceiling();
            return generic_shift(a_ar, deltax, deltay);
        }

        private static T[,] generic_shift<T>(T[,] a_ar, int deltax, int deltay)
        {
            T[,] result = new T[a_ar.GetLength(0), a_ar.GetLength(1)];

            for (int x = 0; x < a_ar.GetLength(0); x++)
            {
                int cx = OverlapOverlayCorrector.Correct_XY(x + deltax, a_ar.GetLength(0));

                for (int y = 0; y < a_ar.GetLength(1); y++)
                {
                    int cy = OverlapOverlayCorrector.Correct_XY(y + deltay, a_ar.GetLength(1));
                    result[cx, cy] = a_ar[x, y];
                }
            }

            return result;
        }
        #endregion
    }
}
